import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import mplfinance as mpf
from tqdm import tqdm
from datetime import datetime


def insert(df, i, df_add):
    # 指定第i行插入一行数据
    df1 = df.iloc[:i, :]
    df2 = df.iloc[i:, :]
    df_new = pd.concat([df1, df_add, df2], ignore_index=True)
    return df_new


num = 100
df = pd.read_csv(r'./个股数据/sec600089Data.csv')
cellVolume = df.TradVolume.sum() / num

demo2 = np.cumsum(df.TradVolume)
bins = list(range(0,num+1)*cellVolume)
wheredf = pd.DataFrame({'TradeIndex':df.TradeIndex,'TradTime':df.TradTime})
wheredf['Vgroup'] = pd.DataFrame({'Vgroup':pd.cut(demo2,bins, labels = range(1,101), retbins=False, include_lowest = True,right=True).values})
wheredf['cumsum'] = np.cumsum(df.TradVolume)
wheredf['diff'] = wheredf[['Vgroup']].diff()

for i in wheredf[wheredf['diff'] == 1].index:
    #取上一个group最后一个和此group第一个
    cumsum_last = wheredf.loc[i-1]['cumsum']
    cumsum_first = wheredf.loc[i]['cumsum']

    #判断阶梯是否包含在其中
    if wheredf.loc[i-1]['Vgroup']*cellVolume < cumsum_first and wheredf.loc[i-1]['Vgroup']*cellVolume > cumsum_last:
        #如果在其中，则插入复制一行
        wheredf = insert(wheredf,i,pd.DataFrame({'TradeIndex':wheredf.iloc[i]['TradeIndex'],'TradTime':wheredf.iloc[i]['TradTime'],'Vgroup':wheredf.iloc[i-1]['Vgroup'],'cumsum':wheredf.iloc[i]['cumsum'],'diff':np.nan},index = [0]))
for i in range(0,len(wheredf[wheredf['diff'] == 1].index)):
    index = wheredf[wheredf['diff'] == 1].index[i]
    #取上一个group最后一个和此group第一个
    cumsum_last = wheredf.loc[index-1]['cumsum']
    cumsum_first = wheredf.loc[index]['cumsum']
    #判断阶梯是否包含在其中
    if wheredf.loc[index-1]['Vgroup']*cellVolume < cumsum_first :
        if wheredf.loc[index-1]['Vgroup']*cellVolume > cumsum_last: #如果在其中，则插入复制一行
            wheredf = insert(wheredf,index,pd.DataFrame({'TradeIndex':wheredf.iloc[index]['TradeIndex'],'TradTime':wheredf.iloc[index]['TradTime'],'Vgroup':wheredf.iloc[index-1]['Vgroup'],'cumsum':wheredf.iloc[index]['cumsum'],'diff':np.nan},index = [0]))

df1 = pd.merge(wheredf, df[['TradeIndex', 'TradPrice']], on='TradeIndex', how='left')


def myf(df):
    mymax = np.max(df)
    open = df[0]
    mymin = np.min(df)
    close = df[-1]
    return [open, mymax, mymin, close]


data_timespan = df1.groupby('Vgroup')['TradTime'].apply(
    lambda x: datetime.strptime(x.iloc[-1], '%H:%M:%S.%f') - datetime.strptime(x.iloc[0], '%H:%M:%S.%f')).reset_index()
data_timespan['TradTime'] = data_timespan['TradTime'].apply(lambda x: x.seconds)

data_price = df1.groupby('Vgroup').apply(lambda x: myf([i for i in x['TradPrice']])).reset_index()
data_price.columns = ['Date', 'OHLC']
data_price['Open'] = data_price['OHLC'].apply(lambda x: x[0])
data_price['High'] = data_price['OHLC'].apply(lambda x: x[1])
data_price['Low'] = data_price['OHLC'].apply(lambda x: x[2])
data_price['Close'] = data_price['OHLC'].apply(lambda x: x[3])
# data_price['Volume'] = data_timespan[['TradTime']]
data_price.drop(columns='OHLC', inplace=True)

# 绘OHLCT图
mc = mpf.make_marketcolors(
    up="red",  # 上涨K线的颜色
    down="green",  # 下跌K线的颜色
    edge="black",  # 蜡烛图箱体的颜色
    volume="blue",  # 成交量柱子的颜色
    wick="black"  # 蜡烛图影线的颜色
)
# 调用make_mpf_style函数，自定义图表样式
# 函数返回一个字典，查看字典包含的数据，按照需求和规范调整参数
style = mpf.make_mpf_style(base_mpl_style="ggplot", marketcolors=mc)

index1 = pd.date_range("2000-01-01", periods=num, freq="D")
OHLC = data_price[['Open', 'High', 'Low', 'Close']]
OHLC.index = index1
mpf.plot(OHLC, type='candle', style=style)  # 使用candlestick_ohlc绘图
